# -*- coding: utf-8 -*-
from mongoengine import Document
from mongoengine.queryset.queryset import QuerySet
from mongoengine.queryset import QuerySetManager
from pymongo.errors import AutoReconnect, OperationFailure
from retry import retry

class RetryQuerySet(QuerySet):
    '''
        重写QuerySet为无限重试
    '''
    @retry(AutoReconnect, tries=4000, delay=10)
    def count(self, with_limit_and_skip=False):
        return QuerySet.count(self,with_limit_and_skip)

    @retry(AutoReconnect, tries=4000, delay=10)
    @retry(OperationFailure, tries=4000, delay=10)
    def first(self):
        return QuerySet.first(self)
    
class RetryQuerySetManager(QuerySetManager):
    queryset_class = RetryQuerySet
    def __init__(self, queryset_func=None):
        QuerySetManager.__init__(self,queryset_func)
        
    @retry(AutoReconnect, tries=4000, delay=10)
    def __get__(self, instance, owner):
        """Descriptor for instantiating a new QuerySet object when
        Document.objects is accessed.
        """
        if instance is not None:
            # Document object being used rather than a document class
            return self

        # owner is the document that contains the QuerySetManager
        queryset_class = owner._meta.get("queryset_class", self.default)
        queryset = queryset_class(owner, owner._get_collection())
        if self.get_queryset:
            arg_count = self.get_queryset.__code__.co_argcount
            if arg_count == 1:
                queryset = self.get_queryset(queryset)
            elif arg_count == 2:
                queryset = self.get_queryset(owner, queryset)
            else:
                queryset = partial(self.get_queryset, owner, queryset)
        return queryset


class RetryDocument(Document):
    '''
        重写Docmuent为无限重试
    '''
    meta = {
        'abstract':True
    }
    objects = RetryQuerySetManager()

    @retry(AutoReconnect, tries=4000, delay=10)
    def save(
        self,
        force_insert=False,
        validate=True,
        clean=True,
        write_concern=None,
        cascade=None,
        cascade_kwargs=None,
        _refs=None,
        save_condition=None,
        signal_kwargs=None,
        **kwargs
    ):
        return Document.save(self,force_insert,validate,clean,write_concern,cascade,cascade_kwargs,_refs,save_condition,signal_kwargs,**kwargs)
